#include <sys/time.h>
#include <errno.h>
#include <fcntl.h>
#include <unistd.h>
#include <assert.h>
#include <cstdio>
#include <strings.h>

#include "pollmgr.h"

// based on Jeremy Stribling and Alex Pesterev's 6.824 lab project

SelectAIO::SelectAIO() : highfds_(0)
{
	FD_ZERO(&rfds_);
	FD_ZERO(&wfds_);

	assert(pipe(pipefd_) == 0);
	FD_SET(pipefd_[0], &rfds_);
	highfds_ = pipefd_[0];

	int flags = fcntl(pipefd_[0], F_GETFL, NULL);
	flags |= O_NONBLOCK;
	fcntl(pipefd_[0], F_SETFL, flags);

	assert(pthread_mutex_init(&m_, NULL) == 0);
}

SelectAIO::~SelectAIO()
{
	assert(pthread_mutex_destroy(&m_) == 0);
}

bool SelectAIO::watch_fd(int fd, poll_flag flag)
{
    pthread_mutex_lock(&m_);
	if (highfds_ <= fd) 
		highfds_ = fd;

	if (flag == PM_RDONLY) {
		FD_SET(fd, &rfds_);
	}else if (flag == PM_WRONLY) {
		FD_SET(fd, &wfds_);
	}else {
		FD_SET(fd, &rfds_);
		FD_SET(fd, &wfds_);
	}

	char tmp = 1;
	assert(write(pipefd_[1], &tmp, sizeof(tmp)) == 1);

    pthread_mutex_unlock(&m_);
    return true;
}

bool SelectAIO::is_watched(int fd, poll_flag flag)
{
    pthread_mutex_lock(&m_);
	if (flag == PM_RDONLY) {
		return FD_ISSET(fd, &rfds_);
	}else if (flag == PM_WRONLY) {
		return FD_ISSET(fd, &wfds_);
	}else{
		return (FD_ISSET(fd, &rfds_) && FD_ISSET(fd, &wfds_));
	}
    pthread_mutex_unlock(&m_);
}

bool SelectAIO::unwatch_fd(int fd, poll_flag flag)
{
    pthread_mutex_lock(&m_);
	if (flag == PM_RDONLY) {
		FD_CLR(fd, &rfds_);
	}else if (flag == PM_WRONLY) {
		FD_CLR(fd, &wfds_);
	}else if (flag == PM_RDWR) {
		FD_CLR(fd, &wfds_);
		FD_CLR(fd, &rfds_);
	}else{
		assert(0);
	}

	if (!FD_ISSET(fd, &rfds_) && !FD_ISSET(fd, &wfds_)) {
		if (fd == highfds_) {
			int newh = pipefd_[0];
			for (int i = 0; i <= highfds_; i++) {
				if (FD_ISSET(i, &rfds_)) {
					newh = i;
				}else if (FD_ISSET(i, &wfds_)) {
					newh = i;
				}
			}
			highfds_ = newh;
		}
	}
	if (flag == PM_RDWR) {
		char tmp = 1;
		assert(write(pipefd_[1], &tmp, sizeof(tmp))==1);
	}
    pthread_mutex_unlock(&m_);
	return (!FD_ISSET(fd, &rfds_) && !FD_ISSET(fd, &wfds_));
}

void
SelectAIO::wait_ready(std::vector<int> &readable, std::vector<int> &writable)
{
	fd_set trfds, twfds;
	int high;

	{
        pthread_mutex_lock(&m_);
		trfds = rfds_;
		twfds = wfds_;
		high = highfds_;
        pthread_mutex_unlock(&m_);
	}

	int ret = select(high + 1, &trfds, &twfds, NULL, NULL);

	if (ret < 0) {
		if (errno == EINTR) {
			return;
		} else {
			perror("select:");
			assert(0);
		}
	}

	for (int fd = 0; fd <= high; fd++) {
		if (fd == pipefd_[0] && FD_ISSET(fd, &trfds)) {
			char tmp;
			assert (read(pipefd_[0], &tmp, sizeof(tmp)) == 1);
			assert(tmp==1);
		}else {
			if (FD_ISSET(fd, &twfds)) {
				writable.push_back(fd);
			}
			if (FD_ISSET(fd, &trfds)) {
				readable.push_back(fd);
			}
		}
	}
}

#ifdef __linux__ 

EPollAIO::EPollAIO()
{
	pollfd_ = epoll_create(MAX_POLL_FDS);
	assert(pollfd_ >= 0);
	bzero(fdstatus_, sizeof(int)*MAX_POLL_FDS);
}

EPollAIO::~EPollAIO()
{
	close(pollfd_);
}

static inline int poll_flag_to_event(poll_flag flag)
{
	int f;
	if (flag == PM_RDONLY) {
		f = EPOLLIN;
	}else if (flag == PM_WRONLY) {
		f = EPOLLOUT;
	}else { //flag == PM_RDWR
		f = EPOLLIN | EPOLLOUT;
	}
	return f;
}

bool EPollAIO::watch_fd(int fd, poll_flag flag)
{
	assert(fd < MAX_POLL_FDS);

	struct epoll_event ev;
	int op = fdstatus_[fd] ? EPOLL_CTL_MOD : EPOLL_CTL_ADD;
	fdstatus_[fd] |= (int)flag;

	ev.events = 0;
	ev.data.fd = fd;

	if (fdstatus_[fd] & PM_RDONLY) {
		ev.events |= EPOLLIN;
	}
	if (fdstatus_[fd] & PM_WRONLY) {
		ev.events |= EPOLLOUT;
	}

	if (flag == PM_RDWR) {
		assert(ev.events == (uint32_t)( EPOLLIN | EPOLLOUT));
	}

	assert(epoll_ctl(pollfd_, op, fd, &ev) == 0);
    return true;
}

bool 
EPollAIO::unwatch_fd(int fd, poll_flag flag)
{
	assert(fd < MAX_POLL_FDS);
	fdstatus_[fd] &= ~(int)flag;

	struct epoll_event ev;
	int op = fdstatus_[fd]? EPOLL_CTL_MOD : EPOLL_CTL_DEL;

	ev.events = EPOLLET;
	ev.data.fd = fd;

	if (fdstatus_[fd] & PM_RDONLY) {
		ev.events |= EPOLLIN;
	}
	if (fdstatus_[fd] & PM_WRONLY) {
		ev.events |= EPOLLOUT;
	}

	if (flag == PM_RDWR) {
		assert(op == EPOLL_CTL_DEL);
	}

	assert(epoll_ctl(pollfd_, op, fd, &ev) == 0);
	return (op == EPOLL_CTL_DEL);
}

bool
EPollAIO::is_watched(int fd, poll_flag flag)
{
	assert(fd < MAX_POLL_FDS);
	return ((fdstatus_[fd] & PM_MASK) == flag);
}

void
EPollAIO::wait_ready(std::vector<int> &readable, std::vector<int> &writable)
{
	int nfds = epoll_wait(pollfd_, ready_,	MAX_POLL_FDS, -1);
	for (int i = 0; i < nfds; i++) {
		if (ready_[i].events & EPOLLIN) {
			readable.push_back(ready_[i].data.fd);
		}
		if (ready_[i].events & EPOLLOUT) {
			writable.push_back(ready_[i].data.fd);
		}
	}
}

#endif
